# -*- coding: utf-8 -*-
import matplotlib.pyplot as plt
import matplotlib.image as img
import numpy as np
import random



def splitBlock(img,x,y):
    # 图片分块
    l = []
    for i in range(img.shape[0]//x):
        for j in range(img.shape[1]//y):
            l.append(img[i*x:(i+1)*x,j*y:(j+1)*y])
    return l


def fm(block, m):
    # 翻转
    val = block.flatten()
    for i in range(len(m)):
        if m[i] == 1:
            # 正翻转
            val[i] = val[i]^1
        elif m[i] == -1:
            # 负翻转
            val[i] = val[i]-1 + (val[i]%2)*2
    return val


def smoothness(img, f=0):
    #平滑度函数
    img = img.flatten()
    for i in range(1, len(img)):
        f += abs(int(img[i])-int(img[i-1]))
    return f



def rs(img, x, y, M, M_):
    img=img.copy()
    RM=0
    SM=0
    RM_=0
    SM_=0
    blocks = splitBlock(img,x,y)
    for block in blocks:
        s0 = smoothness(block)
        s1 = smoothness(fm(block, M)) # 非负翻转
        s2 = smoothness(fm(block, M_))# 非正翻转
        if s1 > s0:
            RM += 1
        elif s1 < s0:
            SM += 1
        if s2 > s0:
            RM_ += 1
        elif s2 < s0:
            SM_ += 1
    
    return RM,SM,RM_,SM_

def m(x,y):
    M = [0]*(x*y//2)+[1]*(x*y-x*y//2)
    np.random.shuffle(M)
    M_=M.copy()
    for i in range(len(M_)):
        if M_[i] == 1:
            M_[i] = -1

    return M,M_


def LSB(img,rate=1):
    # 对图片进行隐写率为rate的LSB隐写
    img = img.copy()
    x,y=img.shape
    a = int(x*y*rate)
    for i in range(a):
        if random.randint(0,1) == 1:
            img[i//x][i%x] = img[i//x][i%x]^1
    return img


def stegRate(img,x,y):

    M, M_ = m(x,y)
    RM0,SM0,RM_0,SM_0 = rs(img,x,y,M,M_)
    RS0 = [RM0,SM0,RM_0,SM_0]
    #print(f"RM={RM0}SM={SM0}\nRM-={RM_0},SM-={SM_0}")

    # 图片最低位取反
    img_ = img.copy()
    for i in range(img_.shape[0]):
        for j in range(img_.shape[1]):
            img_[i,j] = img_[i,j]^1

    RM1,SM1,RM_1,SM_1 = rs(img_,x,y,M,M_)
    RS1 = [RM1,SM1,RM_1,SM_1]
    #print(f"最低位取反后\nRM={RM1}SM={SM1}\nRM-={RM_1},SM-={SM_1}")

    # 计算隐写率
    d0=RM0-SM0
    d1=RM1-SM1
    d2=RM_0-SM_0
    d3=RM_1-SM_1
    args=[2*(d1+d0),d2-d3 -d1-3*d0,d0-d2]
    a,b=np.roots(args)
    a=np.real(a)
    b=np.real(b)
    if a>b:
        rate = b/(b-0.5)
    else:
        rate = a/(a-0.5)
    if rate < 0:
        rate = 0
    print("估计隐写率为：",rate)
    return rate, RS0, RS1

